import math, time
import pycuda.gpuarray
import pycuda.driver as cuda
import pycuda.autoinit

from read_mnist import *
from kernels import sigmoid
import numpy

class SigmoidTest:
  def __init__( self,
                imgs_fname='datasets/train-images.idx3-ubyte',
                labels_fname='datasets/train-labels.idx1-ubyte',
                splits=16,
                mu=0,
                sigma=1
                ):
    self.imgs = read_mnist_images(imgs_fname)
    if labels_fname:
      self.labels = read_mnist_labels(labels_fname)

    self.splits = splits
    self.mu = mu
    self.sigma = sigma

  def gpu(self):
    start = cuda.Event()
    end = cuda.Event()

    gpu_arr = pycuda.gpuarray.GPUArray((len(self.imgs)/self.splits,28,28), numpy.float32)
    gpu_out_arr = pycuda.gpuarray.empty_like(gpu_arr)

    start.record()
    for subset in range(0, self.splits):
      gpu_arr.set(self.imgs[
        (subset*len(self.imgs)/self.splits):((subset+1)*len(self.imgs)/self.splits)
        ].astype(numpy.float32))

      sigmoid(gpu_arr, self.mu, self.sigma, gpu_out_arr)

    end.record()
    end.synchronize()
    
    secs = start.time_till(end)*1e-3
    return secs

  def cpu(self):
    start = time.time()
    the_exp = (self.imgs - self.mu) * self.sigma
    ans = 1/(1+numpy.exp(the_exp))
    secs = time.time() - start
    #print ans[1]

    return secs

  def python(self):
    start = time.time()

    sigmoid = lambda x, mu, sigma: 1/(1+math.exp((x-mu)*sigma))

    ans = [sigmoid(x, self.mu, self.sigma) for x in self.imgs.flat]
    secs = time.time() - start
    #print ans[1]

    return secs

class GpuTest:
  def __init__(self):
    self.start = cuda.Event()
    self.end = cuda.Event()

  def sigmoid_test(self, input, output):
    self.start.record()
    
    sigmoid(input, 0, 1, output)
    
    self.end.record()
    self.end.synchronize()
    
    secs = self.start.time_till(self.end)*1e-3
    return secs

from time import time
class CpuTest:
  def sigmoid_test(self, input, output):
    mu, sigma = (0, 1)
    cpu_nd = input.get()
    start = time()
    the_exp = (cpu_nd - mu) * sigma
    ans = 1/(1+numpy.exp(the_exp))
    secs = time() - start
    print ans[1]

    return secs

import math
class PythonTest:
  def sigmoid_test(self, input, output):
    mu, sigma = (0, 1)
    python_nd = input.get()
    start = time()

    sigmoid = lambda x, mu, sigma: 1/(1+math.exp((x-mu)*sigma))

    ans = [sigmoid(x, 0, 1) for x in python_nd.flat]
    secs = time() - start
    print ans[1]

    return secs
